# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import hashlib
import json
import logging
from typing import List, Optional

from tornado import gen
from tornado.websocket import websocket_connect

from coin.exchange.okex.parse_util import parse_okex_msg
from coin.base.timestamp import get_timestamp


class OkexTradeClientBase(object):
  """
  This class is designed to start and stop only once, after that, the instance
  is abandoned. If you still need the trade client after it is stopped, then you
  have to create a new instance. Use clone method to create an identical new instance.

  There are 3 cases that can cause trade client to stop:
  1. Not receiving pong message for more than 90 seconds.
  2. Websocket is disconnected.
  3. stop method is called.
  """
  def __init__(self,
               api_key: str,
               secret_key: str,
               *,
               symbol_list: Optional[List[str]] = None,
               logger=None):
    self._logger = logger or logging.getLogger(__name__)
    self._api_key = api_key
    self._secret_key = secret_key

    self._ioloop = None
    self._ws = None
    self._stop = False
    self._symbol_list = symbol_list
    self._channels = []
    self._ws_url = None

    self.ready_callback = None
    self.close_callback = None
    self.message_callback = None

    # pong check
    self._last_pong_timestamp = 0

  def _check_pong_timestamp(self):
    time_elapsed = (get_timestamp() - self._last_pong_timestamp) / 1e9
    if time_elapsed > 90:
      self._logger.error('No pong received for %d seconds, restart trade client soon.',
                         time_elapsed)
      self.stop()

  def _sign(self, param: dict):
    req_param = []
    for key in sorted(param.keys()):
      req_param.append('%s=%s' % (str(key), str(param[key])))
    req_param.append('secret_key=%s' % self._secret_key)
    req_str = '&'.join(req_param)
    return hashlib.md5(req_str.encode('utf-8')).hexdigest().upper()

  @gen.coroutine
  def _login(self):
    for i in range(50):
      param = {'api_key': self._api_key}
      param['sign'] = self._sign(param)
      yield self._ws.write_message(json.dumps({'event': 'login', 'parameters': param}))

      msg = yield self._ws.read_message()
      assert msg is not None, 'websocket disconnected.'
      rsp = parse_okex_msg(msg, self._logger)
      assert rsp is not None
      assert rsp['channel'] == 'login'
      if rsp['data']['result']:
        self._logger.info('Login successful!')
        return

      self._logger.error('Login failure. Try again in 5 seconds. msg=%s', rsp)
      yield gen.sleep(5)

  @gen.coroutine
  def _subscribe(self, channel):
    param = {'api_key': self._api_key}
    param['sign'] = self._sign(param)
    yield self._ws.write_message(
        json.dumps({
            'event': 'addChannel', 'channel': channel, 'parameters': param
        }))
    for _ in range(50):
      msg = yield self._ws.read_message()
      assert msg is not None, 'websocket disconnected.'
      rsp = parse_okex_msg(msg, self._logger)
      assert rsp is not None
      if rsp.get('channel', None) == 'addChannel':
        if rsp['data']['result']:
          break
    assert rsp['channel'] == 'addChannel', rsp['channel']
    assert rsp['data']['result'] is True, rsp['data']['result']
    assert rsp['data']['channel'] == channel, rsp['data']['channel']
    self._logger.info('Channel "%s" subscribed' % channel)

  @gen.coroutine
  def _ping(self):
    while not self._stop:
      yield gen.sleep(30)
      if self._ws is not None:
        self._logger.info('ping')
        try:
          yield self._ws.write_message('{"event":"ping"}')
          self._check_pong_timestamp()
        except Exception:
          if self._stop:
            self._logger.info('Trade client (%d) is stopped already!', id(self))
          else:
            raise

  @gen.coroutine
  def _connect(self):
    while True:
      try:
        self._logger.info('OkexTradeClient try to connect to the websocket.')
        self._ws = yield websocket_connect(self._ws_url)
        yield self._login()
        for channel in self._channels:
          yield self._subscribe(channel)
        break
      except Exception:
        self._logger.exception('Connection failed. Reconnecting in 10 seconds...')
        self._close_ws()
        yield gen.sleep(10)

    self._logger.info('Subscription done.')
    if self.ready_callback is not None:
      self.ready_callback()
    self._ioloop.add_callback(self._run)

  @gen.coroutine
  def _run(self):
    while not self._stop:
      self._run_loop_stopped = False
      msg = None
      try:
        msg = yield self._ws.read_message()
      except Exception:
        self._logger.exception('Error reading ws message')
        self._close_ws()

      if msg is None:
        self._logger.error('Trade websocket is closed!')
        self.stop()
        return

      if msg == "":
        yield gen.sleep(0.01)
        continue
      rsp = parse_okex_msg(msg, self._logger)
      if rsp is None:
        yield gen.sleep(0.01)
        continue

      self._dispatch_message(rsp)

  def _dispatch_message(self, msg):
    if self._stop is True:
      return

    if msg.get('event', None) == 'pong':
      self._logger.info('pong')
      self._last_pong_timestamp = get_timestamp()
      return

    try:
      if self.message_callback is not None:
        self.message_callback(msg)
    except Exception:
      self._logger.exception('Process msg failed, %s', msg)

  def _close_ws(self):
    if self._ws is not None:
      self._ws.close()

  def start(self, ioloop):
    if self._stop is True:
      self._logger.error('Okex trade client start method can only be called once.')
      return

    self._ioloop = ioloop
    self._ioloop.add_callback(self._connect)
    self._ioloop.add_callback(self._ping)
    self._last_pong_timestamp = get_timestamp()

  def stop(self, trigger_close_callback=True):
    """Note that even after stop is called, the run method may still running the last loop."""
    if self._stop is True:
      return

    self._stop = True
    self._close_ws()
    if trigger_close_callback and self.close_callback is not None:
      self.close_callback()

  def clone(self):
    trade_client = type(self)(self._api_key,
                              self._secret_key,
                              symbol_list=self._symbol_list,
                              logger=self._logger)
    trade_client.ready_callback = self.ready_callback
    trade_client.close_callback = self.close_callback
    trade_client.message_callback = self.message_callback
    return trade_client
